//#ifdef _ENCRYPT

#include "Openssl.h"
#include <stdlib.h>
#include <string.h>
#include <assert.h>

#define RSA_E_LEN_BYTE 4
#define RSA_N_LEN_BYTE 128
#define RSA_E 65537
#define RSA_PADDING_LEN 11
#define AES_BLOCK_SIZE 16

#define AES_KEY_LEN 128

static EncryptAlgType g_encrypt_alg_type = kTypeAes;
static int rsaInit = 0;
static mbedtls_rsa_context g_rsa;
static mbedtls_aes_context aesctx;
static unsigned char key[BUFFER_SIZE] = {0};

void edp_printf(const char *format, ...)
{
    /*char uart_buf[256];
    va_list ap;

    va_start(ap, format);
    vsnprintf(uart_buf, sizeof(uart_buf), format, ap);
    va_end(ap);

    SXS_TRACE(_AT | TNB_ARG(0) | TSTDOUT, uart_buf);*/
}

/*
 * buffer按十六进制输出
 */
/*void hexdump(const unsigned char *buf, uint32 num)
{
    edp_printf("begin hexdump");

    uint32 i = 0;
    for (; i < num; i++)
    {
        edp_printf("%02X ", buf[i]);
        if ((i + 1) % 8 == 0)
            edp_printf("\n");
    }
    edp_printf("\n");
}*/

static int myrand(void *rng_state, unsigned char *output, size_t len)
{
    size_t use_len;
    int rnd;

    if (rng_state != NULL)
        rng_state = NULL;

    while (len > 0)
    {
        use_len = len;
        if (use_len > sizeof(int))
            use_len = sizeof(int);

        rnd = rand();
        memcpy(output, &rnd, use_len);
        output += use_len;
        len -= use_len;
    }

    return (0);
}

EdpPacket *PacketEncryptReq(EncryptAlgType type)
{
    //int len = 0;
    int ret = 0;
    EdpPacket *pkg = NULL;
    unsigned remainlen = 0;
    unsigned char tmp[RSA_E_LEN_BYTE] = {0};
    if (type != kTypeAes)
    {
        return NULL;
    }
    /* init rsa */
    if (rsaInit == 0)
    {
        mbedtls_rsa_init(&g_rsa, MBEDTLS_RSA_PKCS_V15, 0);
        // gen key pair test
        if (0 != mbedtls_rsa_gen_key(&g_rsa, myrand, NULL, RSA_N_LEN_BYTE * 8, RSA_E))
        { //requrement of ofo :1024bit public key
            edp_printf("mbedtls_rsa_gen_key failed");
            rsaInit = 0;
            return NULL;
        }
        printf("trytrytry:mbedtls_rsa_gen_key success");
        edp_printf("mbedtls_rsa_gen_key success");
        rsaInit = 1;
    }
    pkg = NewBuffer();
    if (pkg == NULL)
        return NULL;
    remainlen = RSA_N_LEN_BYTE + RSA_E_LEN_BYTE + 1;
    WriteByte(pkg, ENCRYPTREQ);
    WriteRemainlen(pkg, remainlen);

    /* write e */
    memset(pkg->_data + pkg->_write_pos, 0, RSA_E_LEN_BYTE);
    //len = BN_bn2bin(g_rsa->e, tmp);
    ret = mbedtls_mpi_write_binary(&g_rsa.E,
                                   tmp, RSA_E_LEN_BYTE);
    if (ret < 0)
    {
        edp_printf("mbedtls_mpi_write_string failed ret = %d", ret);
        DeleteBuffer(&pkg);
        return NULL;
    }
    memcpy(pkg->_data + pkg->_write_pos,
           tmp, RSA_E_LEN_BYTE);
    pkg->_write_pos += RSA_E_LEN_BYTE;

    /* write n */
    //len = BN_bn2bin(g_rsa->N, pkg->_data + pkg->_write_pos);
    ret = mbedtls_mpi_write_binary(&g_rsa.N,
                                   pkg->_data + pkg->_write_pos, RSA_N_LEN_BYTE);
    if (ret < 0)
    {
        edp_printf("mbedtls_mpi_write_string N failed,ret = %d", ret);
        DeleteBuffer(&pkg);
        return NULL;
    }
    //memcpy(pkg->_data + pkg->_write_pos, );
    pkg->_write_pos += RSA_N_LEN_BYTE;

    /* 
     * The type of symmetric encryption algorithm.
     * Right now, only support AES whose code is 1
     */
    g_encrypt_alg_type = type;
    WriteByte(pkg, g_encrypt_alg_type);

    return pkg;
}

int32 UnpackEncryptResp(EdpPacket *pkg)
{
    uint32 remainlen = 0;
    uint16 key_len = 0;

    //unsigned rsa_size = 0;
    //unsigned slice_size = 0;

    //unsigned decrypt_len = 0;
    size_t len = 0;
    unsigned char *from = 0;
    //unsigned char *to = 0;
    //uint32 i = 0;
    //unsigned char rsa_decrypted[10240] = {0};

    mbedtls_aes_init(&aesctx);

    if (ReadRemainlen(pkg, &remainlen))
        return ERR_UNPACK_ENCRYPT_RESP;

    if (ReadUint16(pkg, &key_len))
        return ERR_UNPACK_ENCRYPT_RESP;

    if (remainlen != key_len + 2)
        return ERR_UNPACK_ENCRYPT_RESP;

    //rsa_size = RSA_size(g_rsa);
    //rsa_size = mbedtls_mpi_size(&g_rsa.N);
    //slice_size = rsa_size - RSA_PADDING_LEN;

    from = pkg->_data + pkg->_read_pos;
    //to = key;

    /*for (i=0; i<key_len; i+=rsa_size){
	len = RSA_private_decrypt(rsa_size, from+i, to, 
				  g_rsa, RSA_PKCS1_PADDING);
	
	decrypt_len += len;
	if (decrypt_len > BUFFER_SIZE){
	    return ERR_UNPACK_ENCRYPT_RESP;
	}
    
	to += len;
    }*/
    if (mbedtls_rsa_pkcs1_decrypt(&g_rsa, myrand, NULL, MBEDTLS_RSA_PRIVATE, &len,
                                  from, key,
                                  BUFFER_SIZE) != 0)
    {

        edp_printf("failed\n");

        return (-1);
    }
    if (AES_KEY_LEN != len * 8)
    {
        edp_printf("key length error,len =%d\n", len);
        return -1;
    }
    switch (g_encrypt_alg_type)
    {
    case kTypeAes:
        //if(AES_set_encrypt_key(key, AES_KEY_LEN, &g_aes_encrypt_key) < 0) {
        if (mbedtls_aes_setkey_enc(&aesctx, key, AES_KEY_LEN) < 0)
        {
            return ERR_UNPACK_ENCRYPT_RESP;
        }

        //if(AES_set_decrypt_key(key, AES_KEY_LEN, &g_aes_decrypt_key) < 0) {
        if (mbedtls_aes_setkey_dec(&aesctx, key, AES_KEY_LEN) < 0)
        {
            return ERR_UNPACK_ENCRYPT_RESP;
        }
        break;

    default:
        return ERR_UNPACK_ENCRYPT_RESP;
        break;
    }

    return 0;
}

/*************** AES ***************/
static int aes_encrypt(EdpPacket *pkg, int remain_pos)
{
    uint32 in_len = 0;
    unsigned char *in = NULL;
    unsigned char *out = NULL;
    size_t div = 0;
    size_t mod = 0;
    size_t adder = 0;
    size_t block = 0;
    size_t padding_len = 0;
    unsigned char *padding_addr = NULL;
    size_t len_aft_enc = 0;
    uint8 tmp_buf[5] = {0};
    EdpPacket tmp_pkg;
    int diff = 0;
    int i = 0;

    /* 加密 */
    in_len = pkg->_write_pos - remain_pos;
    in = pkg->_data + remain_pos;
    /* AES 支持加密后数据存放在加密前的buffer中 */
    out = in;

    div = in_len / AES_BLOCK_SIZE;
    mod = in_len % AES_BLOCK_SIZE;

    padding_len = AES_BLOCK_SIZE - mod;
    padding_addr = in + div * AES_BLOCK_SIZE;
    if (mod)
    {
        padding_addr += mod;
    }
    ++div;

    //system("pause");
    /* 填充
     * (1) 如果被加密数据长度刚好是AES_BLOCK_SIZE的整数倍，则在后面
     *  填充AES_BLOCK_SIZE个'0' + AES_BLOCK_SIZE，如果AES_BLOCK_SIZE等于
     *  16，则填充16个 '0'+16
     * (2) 如果不是整数倍，假设最后一段长度为n，则在末尾填充AES_BLOCK_SIZE-n
     * 个'0' + AES_BLOCK_SIZE - n。比如AES_BLOCK_SIZE = 16, n=11, 则在最后
     * 填充5个 '0'+5.
     */

    memset(padding_addr, '0' + padding_len, padding_len); //

    mbedtls_aes_setkey_enc(&aesctx, key, AES_KEY_LEN);
    /* AES 支持加解密前后存放在同一buffer中 */
    for (block = 0; block <= div; ++block)
    {
        adder = block * AES_BLOCK_SIZE;
        //AES_encrypt(in + adder, out + adder, &g_aes_encrypt_key);
        mbedtls_aes_crypt_ecb(&aesctx, MBEDTLS_AES_ENCRYPT, in + adder, out + adder);
        //mbedtls_aes_encrypt(&aesctx, in + adder, out + adder);
    }

    /* 加密后的remainlen会变得大，其占用的空间可能会变大
     * 利用一个临时的EdpPacket来测试加密后remainlen的长度是否发生改变
     * 若改变，则加密后的数据应该依次往后移，以为remainlen留出足够空间
     */
    len_aft_enc = div * AES_BLOCK_SIZE;

    tmp_pkg._data = tmp_buf;
    tmp_pkg._write_pos = 1;
    tmp_pkg._read_pos = 0;
    WriteRemainlen(&tmp_pkg, len_aft_enc);
    diff = tmp_pkg._write_pos - remain_pos;
    if (diff > 0)
    {
        i = len_aft_enc;
        for (; i > 0; i--)
        {
            *(in + i + diff - 1) = *(in + i - 1);
        }
    }

    pkg->_write_pos = 1;
    pkg->_read_pos = 0;
    WriteRemainlen(pkg, len_aft_enc);
    pkg->_write_pos += len_aft_enc;

    return len_aft_enc;
}

static int aes_decrypt(EdpPacket *pkg, int remain_pos)
{
    size_t in_len = 0;
    unsigned char *in = NULL;
    unsigned char *out = NULL;
    size_t offset = 0;
    size_t padding_len = 0;
    uint32 len_aft_dec = 0;
    uint8 tmp_buf[5] = {0};
    EdpPacket tmp_pkg;
    int diff = 0;
    int i = 0;

    in_len = pkg->_write_pos - pkg->_read_pos;
    in = pkg->_data + pkg->_read_pos;
    out = in;
    mbedtls_aes_setkey_dec(&aesctx, key, AES_KEY_LEN);
    for (offset = 0; (offset + AES_BLOCK_SIZE) <= in_len; offset += AES_BLOCK_SIZE)
    {
        mbedtls_aes_decrypt(&aesctx, in + offset, out + offset);
        //AES_decrypt(in + offset, out + offset, &g_aes_decrypt_key);
        //mbedtls_aes_crypt_ecb(&aesctx,MBEDTLS_AES_DECRYPT,in+offset,out+offset);
    }

    padding_len = *(in + offset - 1) - '0';
    if (padding_len > AES_BLOCK_SIZE)
    {
        return -1;
    }

    /* 解密后的remainlen会变小，其占用空间可能变小
     * 利用一个临时的EdpPacket来测试加密后remainlen的长度是否发生改变
     * 若改变，则解密后的数据应该依次往前移，以消除多余空间
     */
    len_aft_dec = offset - padding_len;
    tmp_pkg._data = tmp_buf;
    tmp_pkg._write_pos = 1;
    tmp_pkg._read_pos = 0;
    WriteRemainlen(&tmp_pkg, len_aft_dec);

    diff = remain_pos - tmp_pkg._write_pos;
    if (diff > 0)
    {
        i = 0;
        for (i = 0; i < len_aft_dec; i++)
        {
            *(in + i - diff) = *(in + i);
        }
    }

    pkg->_write_pos = 1;
    pkg->_read_pos = 0;
    WriteRemainlen(pkg, len_aft_dec);
    pkg->_read_pos = 1;
    pkg->_write_pos += len_aft_dec;

    return len_aft_dec;
}

/*************** end AES ***************/

int SymmEncrypt(EdpPacket *pkg)
{
    int ret = 0;
    uint32 remain_len = 0;
    uint32 remain_pos = 0;

    pkg->_read_pos = 1;
    ReadRemainlen(pkg, &remain_len);
    assert(remain_len == (pkg->_write_pos - pkg->_read_pos));
    if (remain_len == 0)
    { /* no data for encrypting */
        pkg->_read_pos = 1;
        return 0;
    }
    remain_pos = pkg->_read_pos;

    switch (g_encrypt_alg_type)
    {
    case kTypeAes:
        ret = aes_encrypt(pkg, remain_pos);
        break;

    default:
        ret = -1;
        break;
    }

    return ret;
}

int SymmDecrypt(EdpPacket *pkg)
{
    int ret = 0;
    uint32 remain_len = 0;
    uint32 remain_pos = 0;

    pkg->_read_pos = 1;
    ReadRemainlen(pkg, &remain_len);
    assert(remain_len == (pkg->_write_pos - pkg->_read_pos));
    if (remain_len == 0)
    { /* no data for decrypting */
        pkg->_read_pos = 1;
        return 0;
    }
    remain_pos = pkg->_read_pos;

    switch (g_encrypt_alg_type)
    {
    case kTypeAes:
        ret = aes_decrypt(pkg, remain_pos);
        break;

    default:
        ret = -1;
        break;
    }

    return ret;
}
//#endif
